package cn.bugstack.chatgpt.test;

import cn.bugstack.chatgpt.common.Constants;
import cn.bugstack.chatgpt.domain.chat.ChatCompletionRequest;
import cn.bugstack.chatgpt.domain.chat.ChatCompletionResponse;
import cn.bugstack.chatgpt.domain.chat.Message;
import cn.bugstack.chatgpt.domain.completions.ComRequest;
import cn.bugstack.chatgpt.domain.completions.ComResponse;
import cn.bugstack.chatgpt.domain.images.ImageRequest;
import cn.bugstack.chatgpt.domain.images.ImageResponse;
import cn.bugstack.chatgpt.domain.qa.QACompletionResponse;
import cn.bugstack.chatgpt.domain.web.WebRequest;
import cn.bugstack.chatgpt.domain.web.WebResponse;
import cn.bugstack.chatgpt.domain.web.WebResult;
import cn.bugstack.chatgpt.session.Configuration;
import cn.bugstack.chatgpt.session.OpenAiSession;
import cn.bugstack.chatgpt.session.OpenAiSessionFactory;
import cn.bugstack.chatgpt.session.defaults.DefaultOpenAiSessionFactory;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.junit.Before;
import org.junit.Test;

import java.util.Collections;
import java.util.concurrent.CountDownLatch;

/**
 * @author 小傅哥，微信：fustack
 * @description 单元测试
 * @github https://github.com/fuzhengwei
 * @Copyright 公众号：bugstack虫洞栈 | 博客：https://bugstack.cn - 沉淀、分享、成长，让自己和他人都能有所收获！
 */
@Slf4j
public class ApiTest {

    private OpenAiSession openAiSession;

    @Before
    public void test_OpenAiSessionFactory() {
        // 1. 配置文件【如果你从小傅哥获取key会给你提供apihost，你可以分别替换下】
        Configuration configuration = new Configuration();
        configuration.setApiHost("https://open.bigmodel.cn/api/paas/");
        configuration.setApiKey("da9dad54d0884d5db6f6001a45b63a83.C3J5shVPSHb8kV8F");
        // 可以根据课程首页评论置顶说明获取 apihost、apikey；https://t.zsxq.com/0d3o5FKvc
        //configuration.setAuthToken("eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ4ZmciLCJleHAiOjE2ODMyODE2NzEsImlhdCI6MTY4MzI3ODA3MSwianRpIjoiMWUzZTkwYjYtY2UyNy00NzNlLTk5ZTYtYWQzMWU1MGVkNWE4IiwidXNlcm5hbWUiOiJ4ZmcifQ.YgQRJ2U5-9uydtd6Wbkg2YatsoX-y8mS_OJ3FdNRaX0");
        // 2. 会话工厂
        OpenAiSessionFactory factory = new DefaultOpenAiSessionFactory(configuration);
        // 3. 开启会话
        this.openAiSession = factory.openSession();
    }

    /**
     * 简单问答模式，方法废弃推荐使用 test_chat_completions  找不到glm问答模型的host应该没有，找不到对应模型名称
     */
    @Test
    public void test_qa_completions() throws JsonProcessingException {
        QACompletionResponse response01 = openAiSession.completions("写个java冒泡排序");
        log.info("测试结果：{}", new ObjectMapper().writeValueAsString(response01.getChoices()));
    }

    /**
     * 此对话模型 3.5 接近于官网体验
     *
     * 文档：https://platform.openai.com/docs/guides/text-generation/chat-completions-api
     * 你可以替换能访问的 apihost【https://api.openai.com】 和 $OPENAI_API_KEY 进行 http 测试
     * curl https://api.openai.com/v1/chat/completions \
     *   -H "Content-Type: application/json" \
     *   -H "Authorization: Bearer $OPENAI_API_KEY" \
     *   -d '{
     *     "model": "gpt-3.5-turbo",
     *     "messages": [
     *       {
     *         "role": "system",
     *         "content": "You are a helpful assistant."
     *       },
     *       {
     *         "role": "user",
     *         "content": "Who won the world series in 2020?"
     *       },
     *       {
     *         "role": "assistant",
     *         "content": "The Los Angeles Dodgers won the World Series in 2020."
     *       },
     *       {
     *         "role": "user",
     *         "content": "Where was it played?"
     *       }
     *     ]
     *   }'
     */
    @Test
    public void test_chat_completions() {
        // 1. 创建参数
        ChatCompletionRequest chatCompletion = ChatCompletionRequest
                .builder()
                .messages(Collections.singletonList(Message.builder().role(Constants.Role.USER).content("1+1").build()))
                .model(ChatCompletionRequest.Model.GLM_4_FLASH.getCode())
                .build();
        // 2. 发起请求
        ChatCompletionResponse chatCompletionResponse = openAiSession.completions(chatCompletion);
        // 3. 解析结果
        chatCompletionResponse.getChoices().forEach(e -> {
            log.info("测试结果：{}", e.getMessage());
        });
    }

    @Test
    public void test_chat_completions_stream() throws JsonProcessingException, InterruptedException {
        // 1. 创建参数
        ChatCompletionRequest chatCompletion = ChatCompletionRequest
                .builder()
                .stream(true)
                .messages(Collections.singletonList(Message.builder().role(Constants.Role.USER).content("1+1").build()))
                .model(ChatCompletionRequest.Model.GLM_4.getCode())
                .build();
        for (int i = 0; i < 1; i++) {
            // 2. 发起请求
            EventSource eventSource = openAiSession.chatCompletions(chatCompletion, new EventSourceListener() {
                @Override
                public void onEvent(EventSource eventSource, String id, String type, String data) {
                    log.info("测试结果：{}", data);
                }
            });
        }
        // 等待
        new CountDownLatch(1).await();
/*        // 2. 发起请求
        ChatCompletionResponse chatCompletionResponse = openAiSession.completions(chatCompletion);
        // 3. 解析结果
        chatCompletionResponse.getChoices().forEach(e -> {
            log.info("测试结果：{}", e.getMessage());
        });*/
    }

    @Test
    public void test_image_generation(){
        // 1. 创建参数
        ImageRequest imageRequest = ImageRequest
                .builder()
                .prompt("一只可爱的萨摩耶")
                .model(ImageRequest.Model.COGVIEW_3_FLASH.getCode())
                .build();
        // 2. 发起请求
        ImageResponse imageResponse = openAiSession.imageGenerations(imageRequest);
        // 3. 解析结果
        imageResponse.getData().forEach(e -> {
            log.info("测试结果：{}", e.getUrl());
        });
    }


    @Test
    public void test_web_search(){
        // 1. 创建参数
        WebRequest webRequest = WebRequest
                .builder()
                .searchQuery("差分隐私在图数据中的应用")
                .searchEngine(WebRequest.Model.SEARCH_PRO.getCode())
                .searchIntent(true)
                .build();
        // 2. 发起请求
        WebResponse webResponse = openAiSession.webSearch(webRequest);
        // 3. 解析结果
        /*webResponse.getSearchResult().forEach(e -> {
            log.info("测试结果：{}", e);
        });*/
        if (webResponse.getSearchResult() != null) {
            for (WebResult result : webResponse.getSearchResult()) {
                System.out.println("【标题】：" + result.getTitle());
                System.out.println("【摘要】：" + result.getContent());
                System.out.println("【链接】：" + result.getLink());
                System.out.println("【来源】：" + result.getMedia());
                System.out.println("--------------------------------------------------");
            }
        } else {
            System.out.println("⚠️ 未返回任何搜索结果。");
        }
    }

    @Test
    public void test_chat_completions2() {
        // 1. 创建参数
        ComRequest comRequest = ComRequest
                .builder()
                .messages(Collections.singletonList(Message.builder().role(Constants.Role.USER).content("解释一下什么是MCP（模型上下文协议）").build()))
                .model(ComRequest.Model.GLM_Z1_FLASH.getCode())
                .build();
        // 2. 发起请求
        ComResponse comResponse = openAiSession.chatCom(comRequest);
        // 3. 解析结果
        comResponse.getChoices().forEach(e -> {
            log.info("测试结果：{}", e.getMessage());
        });
    }
}
